import paddle.fluid as fluid
from paddle.fluid.dygraph import Conv2D, Pool2D, Linear

# 下面分别实现了“上采样”和“转置卷积”两种方式实现的生成网络。注释掉其中一个版本可测试另一个。

# 通过上采样扩大特征图的版本
class G(fluid.dygraph.Layer):
    def __init__(self, name_scope):
        super(G, self).__init__(name_scope)
        name_scope = self.full_name()
        # 第一组全连接和BN层
        self.fc1 = Linear(input_dim=100, output_dim=1024)
        self.bn1 = fluid.dygraph.BatchNorm(num_channels=1024, act='tanh')
        # 第二组全连接和BN层
        self.fc2 = Linear(input_dim=1024, output_dim=128 * 7 * 7)
        self.bn2 = fluid.dygraph.BatchNorm(num_channels=128 * 7 * 7, act='tanh')
        # 第一组卷积运算（卷积前进行上采样，以扩大特征图）
        # 注：此处使用转置卷积的效果似乎不如上采样后直接用卷积，转置卷积生成的图片噪点较多
        self.conv1 = Conv2D(num_channels=128, num_filters=64, filter_size=5, padding=2)
        self.bn3 = fluid.dygraph.BatchNorm(num_channels=64, act='tanh')
        # 第二组卷积运算（卷积前进行上采样，以扩大特征图）
        self.conv2 = Conv2D(num_channels=64, num_filters=1, filter_size=5, padding=2, act='tanh')

    def forward(self, z):
        z = fluid.layers.reshape(z, shape=[-1, 100])
        y = self.fc1(z)
        y = self.bn1(y)
        y = self.fc2(y)
        y = self.bn2(y)
        y = fluid.layers.reshape(y, shape=[-1, 128, 7, 7])
        # 第一组卷积前进行上采样以扩大特征图
        y = fluid.layers.image_resize(y, scale=2)
        y = self.conv1(y)
        y = self.bn3(y)
        # 第二组卷积前进行上采样以扩大特征图
        y = fluid.layers.image_resize(y, scale=2)
        y = self.conv2(y)
        return y


# 通过转置卷积扩大特征图的版本
# class G(fluid.dygraph.Layer):
#     def __init__(self, name_scope):
#         super(G, self).__init__(name_scope)
#         name_scope = self.full_name()
#         # 第一组全连接和BN层
#         self.fc1 = Linear(input_dim=100, output_dim=1024)
#         self.bn1 = fluid.dygraph.BatchNorm(num_channels=1024, act='leaky_relu')
#         # 第二组全连接和BN层
#         self.fc2 = Linear(input_dim=1024, output_dim=128*7*7)
#         self.bn2 = fluid.dygraph.BatchNorm(num_channels=128*7*7, act='leaky_relu')
#         # 第一组转置卷积运算
#         self.convtrans1 = Conv2DTranspose(128, 64, 4, stride=2, padding=1)
#         self.bn3 = fluid.dygraph.BatchNorm(64, act='leaky_relu')
#         # 第二组转置卷积运算
#         self.convtrans2 = Conv2DTranspose(64, 1, 4, stride=2, padding=1, act='leaky_relu')

#     def forward(self, z):
#         z = fluid.layers.reshape(z, shape=[-1, 100])
#         y = self.fc1(z)
#         y = self.bn1(y)
#         y = self.fc2(y)
#         y = self.bn2(y)
#         y = fluid.layers.reshape(y, shape=[-1, 128, 7, 7])
#         y = self.convtrans1(y)
#         y = self.bn3(y)
#         y = self.convtrans2(y)
#         return y

class D(fluid.dygraph.Layer):
    def __init__(self, name_scope):
        super(D, self).__init__(name_scope)
        name_scope = self.full_name()
        # 第一组卷积池化
        self.conv1 = Conv2D(num_channels=1, num_filters=64, filter_size=3)
        self.bn1 = fluid.dygraph.BatchNorm(num_channels=64, act='relu')
        self.pool1 = Pool2D(pool_size=2, pool_stride=2)
        # 第二组卷积池化
        self.conv2 = Conv2D(num_channels=64, num_filters=128, filter_size=3)
        self.bn2 = fluid.dygraph.BatchNorm(num_channels=128, act='relu')
        self.pool2 = Pool2D(pool_size=2, pool_stride=2)
        # 全连接输出层
        self.fc1 = Linear(input_dim=128 * 5 * 5, output_dim=1024)
        self.bnfc1 = fluid.dygraph.BatchNorm(num_channels=1024, act='relu')
        self.fc2 = Linear(input_dim=1024, output_dim=1)

    def forward(self, img):
        y = self.conv1(img)
        y = self.bn1(y)
        y = self.pool1(y)
        y = self.conv2(y)
        y = self.bn2(y)
        y = self.pool2(y)
        y = fluid.layers.reshape(y, shape=[-1, 128 * 5 * 5])
        y = self.fc1(y)
        y = self.bnfc1(y)
        y = self.fc2(y)

        return y
